{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, os, pickle\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "human_categories = ['adult', 'child', 'baby']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obj_categories = ['adult', 'car', 'guitar', 'chair', 'handbag', 'toy', 'baby_seat', 'cat', 'bottle', 'backpack', 'motorcycle', 'ball/sports_ball', 'laptop', 'table', 'surfboard', 'camera', 'sofa', 'screen/monitor', 'bicycle', 'vegetables', 'dog', 'fruits', 'cake', 'cellphone', 'cup', 'bench', 'snowboard', 'skateboard', 'bread', 'bus/truck', 'ski', 'suitcase', 'stool', 'bat', 'elephant', 'fish', 'baby_walker', 'dish', 'watercraft', 'scooter', 'pig', 'refrigerator', 'horse', 'crab', 'bird', 'piano', 'cattle/cow', 'lion', 'chicken', 'camel', 'electric_fan', 'toilet', 'sheep/goat', 'rabbit', 'train', 'penguin', 'hamster/rat', 'snake', 'frisbee', 'aircraft', 'oven', 'racket', 'faucet', 'antelope', 'duck', 'stop_sign', 'sink', 'kangaroo', 'stingray', 'turtle', 'tiger', 'crocodile', 'bear', 'microwave', 'traffic_light', 'panda', 'leopard', 'squirrel']\n",
    "with open('obj_categories.json', 'r') as f:\n",
    "    obj_categories = json.load(f)\n",
    "\n",
    "# obj_to_idx = {}\n",
    "# for i, obj in enumerate(obj_categories):\n",
    "#     obj_to_idx[obj] = i\n",
    "# idx_to_obj = {v:k for k, v in obj_to_idx.items()}\n",
    "# print(len(obj_categories))\n",
    "\n",
    "with open('obj_to_idx.pkl', 'rb') as f:\n",
    "    obj_to_idx = pickle.load(f)\n",
    "with open('idx_to_obj.pkl', 'rb') as f:\n",
    "    idx_to_obj = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'person': 0,\n",
       " 'car': 1,\n",
       " 'guitar': 2,\n",
       " 'chair': 3,\n",
       " 'handbag': 4,\n",
       " 'toy': 5,\n",
       " 'baby_seat': 6,\n",
       " 'cat': 7,\n",
       " 'bottle': 8,\n",
       " 'backpack': 9,\n",
       " 'motorcycle': 10,\n",
       " 'ball/sports_ball': 11,\n",
       " 'laptop': 12,\n",
       " 'table': 13,\n",
       " 'surfboard': 14,\n",
       " 'camera': 15,\n",
       " 'sofa': 16,\n",
       " 'screen/monitor': 17,\n",
       " 'bicycle': 18,\n",
       " 'vegetables': 19,\n",
       " 'dog': 20,\n",
       " 'fruits': 21,\n",
       " 'cake': 22,\n",
       " 'cellphone': 23,\n",
       " 'cup': 24,\n",
       " 'bench': 25,\n",
       " 'snowboard': 26,\n",
       " 'skateboard': 27,\n",
       " 'bread': 28,\n",
       " 'bus/truck': 29,\n",
       " 'ski': 30,\n",
       " 'suitcase': 31,\n",
       " 'stool': 32,\n",
       " 'bat': 33,\n",
       " 'elephant': 34,\n",
       " 'fish': 35,\n",
       " 'baby_walker': 36,\n",
       " 'dish': 37,\n",
       " 'watercraft': 38,\n",
       " 'scooter': 39,\n",
       " 'pig': 40,\n",
       " 'refrigerator': 41,\n",
       " 'horse': 42,\n",
       " 'crab': 43,\n",
       " 'bird': 44,\n",
       " 'piano': 45,\n",
       " 'cattle/cow': 46,\n",
       " 'lion': 47,\n",
       " 'chicken': 48,\n",
       " 'camel': 49,\n",
       " 'electric_fan': 50,\n",
       " 'toilet': 51,\n",
       " 'sheep/goat': 52,\n",
       " 'rabbit': 53,\n",
       " 'train': 54,\n",
       " 'penguin': 55,\n",
       " 'hamster/rat': 56,\n",
       " 'snake': 57,\n",
       " 'frisbee': 58,\n",
       " 'aircraft': 59,\n",
       " 'oven': 60,\n",
       " 'racket': 61,\n",
       " 'faucet': 62,\n",
       " 'antelope': 63,\n",
       " 'duck': 64,\n",
       " 'stop_sign': 65,\n",
       " 'sink': 66,\n",
       " 'kangaroo': 67,\n",
       " 'stingray': 68,\n",
       " 'turtle': 69,\n",
       " 'tiger': 70,\n",
       " 'crocodile': 71,\n",
       " 'bear': 72,\n",
       " 'microwave': 73,\n",
       " 'traffic_light': 74,\n",
       " 'panda': 75,\n",
       " 'leopard': 76,\n",
       " 'squirrel': 77}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obj_to_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # preprocessing label order from detected results to my defined label order\n",
    "# import numpy as np\n",
    "\n",
    "# with open('vidvrd-mff/data/vidor/object_labels.txt') as f:\n",
    "#     vidor_classes = f.readlines()\n",
    "\n",
    "# vidor_classes = [cls.strip() for cls in vidor_classes]\n",
    "# print(vidor_classes)\n",
    "\n",
    "# old_idx_to_new = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['person',\n",
       " 'car',\n",
       " 'guitar',\n",
       " 'chair',\n",
       " 'handbag',\n",
       " 'toy',\n",
       " 'baby_seat',\n",
       " 'cat',\n",
       " 'bottle',\n",
       " 'backpack',\n",
       " 'motorcycle',\n",
       " 'ball/sports_ball',\n",
       " 'laptop',\n",
       " 'table',\n",
       " 'surfboard',\n",
       " 'camera',\n",
       " 'sofa',\n",
       " 'screen/monitor',\n",
       " 'bicycle',\n",
       " 'vegetables',\n",
       " 'dog',\n",
       " 'fruits',\n",
       " 'cake',\n",
       " 'cellphone',\n",
       " 'cup',\n",
       " 'bench',\n",
       " 'snowboard',\n",
       " 'skateboard',\n",
       " 'bread',\n",
       " 'bus/truck',\n",
       " 'ski',\n",
       " 'suitcase',\n",
       " 'stool',\n",
       " 'bat',\n",
       " 'elephant',\n",
       " 'fish',\n",
       " 'baby_walker',\n",
       " 'dish',\n",
       " 'watercraft',\n",
       " 'scooter',\n",
       " 'pig',\n",
       " 'refrigerator',\n",
       " 'horse',\n",
       " 'crab',\n",
       " 'bird',\n",
       " 'piano',\n",
       " 'cattle/cow',\n",
       " 'lion',\n",
       " 'chicken',\n",
       " 'camel',\n",
       " 'electric_fan',\n",
       " 'toilet',\n",
       " 'sheep/goat',\n",
       " 'rabbit',\n",
       " 'train',\n",
       " 'penguin',\n",
       " 'hamster/rat',\n",
       " 'snake',\n",
       " 'frisbee',\n",
       " 'aircraft',\n",
       " 'oven',\n",
       " 'racket',\n",
       " 'faucet',\n",
       " 'antelope',\n",
       " 'duck',\n",
       " 'stop_sign',\n",
       " 'sink',\n",
       " 'kangaroo',\n",
       " 'stingray',\n",
       " 'turtle',\n",
       " 'tiger',\n",
       " 'crocodile',\n",
       " 'bear',\n",
       " 'microwave',\n",
       " 'traffic_light',\n",
       " 'panda',\n",
       " 'leopard',\n",
       " 'squirrel']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('obj_categories.json', 'r') as f:\n",
    "    obj_categories = json.load(f)\n",
    "obj_categories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['lean_on',\n",
       " 'watch',\n",
       " 'above',\n",
       " 'next_to',\n",
       " 'behind',\n",
       " 'away',\n",
       " 'towards',\n",
       " 'in_front_of',\n",
       " 'hit',\n",
       " 'hold',\n",
       " 'wave',\n",
       " 'pat',\n",
       " 'carry',\n",
       " 'point_to',\n",
       " 'touch',\n",
       " 'play(instrument)',\n",
       " 'release',\n",
       " 'ride',\n",
       " 'grab',\n",
       " 'lift',\n",
       " 'use',\n",
       " 'press',\n",
       " 'inside',\n",
       " 'caress',\n",
       " 'pull',\n",
       " 'get_on',\n",
       " 'cut',\n",
       " 'hug',\n",
       " 'bite',\n",
       " 'open',\n",
       " 'close',\n",
       " 'throw',\n",
       " 'kick',\n",
       " 'drive',\n",
       " 'get_off',\n",
       " 'push',\n",
       " 'wave_hand_to',\n",
       " 'feed',\n",
       " 'chase',\n",
       " 'kiss',\n",
       " 'speak_to',\n",
       " 'beneath',\n",
       " 'smell',\n",
       " 'clean',\n",
       " 'lick',\n",
       " 'squeeze',\n",
       " 'shake_hand_with',\n",
       " 'knock',\n",
       " 'hold_hand_of',\n",
       " 'shout_at']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('pred_categories.json', 'r') as f:\n",
    "    pred_categories = json.load(f)\n",
    "pred_categories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "50\n",
      "{'lean_on': 0, 'watch': 1, 'above': 2, 'next_to': 3, 'behind': 4, 'away': 5, 'towards': 6, 'in_front_of': 7, 'hit': 8, 'hold': 9, 'wave': 10, 'pat': 11, 'carry': 12, 'point_to': 13, 'touch': 14, 'play(instrument)': 15, 'release': 16, 'ride': 17, 'grab': 18, 'lift': 19, 'use': 20, 'press': 21, 'inside': 22, 'caress': 23, 'pull': 24, 'get_on': 25, 'cut': 26, 'hug': 27, 'bite': 28, 'open': 29, 'close': 30, 'throw': 31, 'kick': 32, 'drive': 33, 'get_off': 34, 'push': 35, 'wave_hand_to': 36, 'feed': 37, 'chase': 38, 'kiss': 39, 'speak_to': 40, 'beneath': 41, 'smell': 42, 'clean': 43, 'lick': 44, 'squeeze': 45, 'shake_hand_with': 46, 'knock': 47, 'hold_hand_of': 48, 'shout_at': 49} {0: 'lean_on', 1: 'watch', 2: 'above', 3: 'next_to', 4: 'behind', 5: 'away', 6: 'towards', 7: 'in_front_of', 8: 'hit', 9: 'hold', 10: 'wave', 11: 'pat', 12: 'carry', 13: 'point_to', 14: 'touch', 15: 'play(instrument)', 16: 'release', 17: 'ride', 18: 'grab', 19: 'lift', 20: 'use', 21: 'press', 22: 'inside', 23: 'caress', 24: 'pull', 25: 'get_on', 26: 'cut', 27: 'hug', 28: 'bite', 29: 'open', 30: 'close', 31: 'throw', 32: 'kick', 33: 'drive', 34: 'get_off', 35: 'push', 36: 'wave_hand_to', 37: 'feed', 38: 'chase', 39: 'kiss', 40: 'speak_to', 41: 'beneath', 42: 'smell', 43: 'clean', 44: 'lick', 45: 'squeeze', 46: 'shake_hand_with', 47: 'knock', 48: 'hold_hand_of', 49: 'shout_at'}\n"
     ]
    }
   ],
   "source": [
    "# pred_categories = [\"lean_on\",\"watch\",\"above\",\"next_to\",\"behind\",\"away\",\"towards\",\"in_front_of\",\"hit\",\"hold\",\"wave\",\"pat\",\"carry\",\"point_to\",\"touch\",\"play(instrument)\",\"release\",\"ride\",\"grab\",\"lift\",\"use\",\"press\",\"inside\",\"caress\",\"pull\",\"get_on\",\"cut\",\"hug\",\"bite\",\"open\",\"close\",\"throw\",\"kick\",\"drive\",\"get_off\",\"push\",\"wave_hand_to\",\"feed\",\"chase\",\"kiss\",\"speak_to\",\"beneath\",\"smell\",\"clean\",\"lick\",\"squeeze\",\"shake_hand_with\",\"knock\",\"hold_hand_of\",\"shout_at\"]\n",
    "# with open('pred_categories.json', 'w') as f:\n",
    "#     json.dump(pred_categories, f)\n",
    "with open('pred_categories.json', 'r') as f:\n",
    "    pred_categories = json.load(f)\n",
    "print(len(pred_categories))\n",
    "\n",
    "# pred_to_idx = {}\n",
    "# for i, pred in enumerate(pred_categories):\n",
    "#     pred_to_idx[pred] = i\n",
    "# idx_to_pred = {v:k for k, v in pred_to_idx.items()}\n",
    "\n",
    "# with open('pred_to_idx.pkl', 'wb') as f:\n",
    "#     pickle.dump(pred_to_idx, f)\n",
    "# with open('idx_to_pred.pkl', 'wb') as f:\n",
    "#     pickle.dump(idx_to_pred, f)\n",
    "\n",
    "with open('pred_to_idx.pkl', 'rb') as f:\n",
    "    pred_to_idx = pickle.load(f)\n",
    "with open('idx_to_pred.pkl', 'rb') as f:\n",
    "    idx_to_pred = pickle.load(f)\n",
    "print(pred_to_idx, idx_to_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 131/131 [00:15<00:00,  8.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nb_empty_img_detection_result: 735\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "def convert_vidor_to_ava_label(annot_dir):\n",
    "    frame_annots = {}\n",
    "    used_video_dict = set()\n",
    "    \n",
    "    val_traj_path = 'vidvrd-mff/data/vidor/vidor_val_object_trajectories.json'\n",
    "    with open(val_traj_path, 'rb') as f:\n",
    "        val_traj = json.load(f)\n",
    "    val_traj = val_traj['results']\n",
    "    \n",
    "    nb_empty_img_detection_result = 0\n",
    "    \n",
    "    for folder in tqdm(os.listdir(annot_dir)):\n",
    "        for video_json in os.listdir(os.path.join(annot_dir, folder)):\n",
    "            with open(os.path.join(annot_dir, folder, video_json), 'r') as f:\n",
    "                annot = json.load(f)\n",
    "            \n",
    "            folder_video_id = folder + '/' + annot['video_id']\n",
    "            predicted_val_trajs = val_traj[folder_video_id] # a list\n",
    "            # keys of a trajectory: \n",
    "            # dict_keys(['category', 'start_fid', 'org_end_fid', 'trajectory', \n",
    "            #            'end_fid', 'height', 'width', 'score', 'org_start_fid'])\n",
    "            \n",
    "            \n",
    "            if abs(annot['fps'] - 29.97) < 0.1:\n",
    "                fps = 30\n",
    "            elif annot['fps'] - 24 < 1.01: # fps 24, 25\n",
    "                fps = 24\n",
    "            else:\n",
    "                import pdb; pdb.set_trace()\n",
    "\n",
    "            for i in range(annot['frame_count']):\n",
    "                if (i - (fps // 2)) % fps != 0: # not middle frame\n",
    "                    continue\n",
    "                \n",
    "                idx = i - 1\n",
    "                # idx = i * fps + fps // 2 - 1\n",
    "                frame_id = folder_video_id + '_' + f'{i // fps + 2:06d}'\n",
    "                \n",
    "                for tid, trajectory in enumerate(predicted_val_trajs):\n",
    "                    if trajectory['start_fid'] <= i < trajectory['end_fid']:\n",
    "                        middle_frame_timestamp = i // fps + 1\n",
    "                        if frame_id not in frame_annots:\n",
    "                            frame_annots[frame_id] = []\n",
    "                        frame_annots[frame_id].append({\n",
    "                            'image_id': frame_id,\n",
    "                            'bbox': trajectory['trajectory'][f'{i:06d}'],\n",
    "                            'category_id': obj_to_idx['person'] if trajectory['category'] in human_categories else obj_to_idx[trajectory['category']],\n",
    "                            'score': trajectory['score'],\n",
    "                            'tid': tid,\n",
    "                        })\n",
    "                        \n",
    "                        used_video_dict.add(folder + '/' + annot['video_id'])\n",
    "                \n",
    "                if frame_id not in frame_annots:\n",
    "                    nb_empty_img_detection_result += 1\n",
    "    \n",
    "    print('nb_empty_img_detection_result:', nb_empty_img_detection_result)\n",
    "    return frame_annots, used_video_dict\n",
    "        \n",
    "# train_annot_dir = 'annotation/training'\n",
    "val_annot_dir = 'annotation/validation'\n",
    "\n",
    "# train_frame_annots, train_video_dict = convert_vidor_to_ava_label(train_annot_dir)\n",
    "det_val_frame_annots, val_video_dict = convert_vidor_to_ava_label(val_annot_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('train_frame_annots.json', 'w') as f:\n",
    "#     json.dump(train_frame_annots, f)\n",
    "with open('det_val_frame_annots.json', 'w') as f:\n",
    "    json.dump(det_val_frame_annots, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('det_val_frame_annots.json', 'r') as f:\n",
    "    val_frame_annots = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28560"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(val_frame_annots)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "with open('val_instances_predictions_train_small_vidor_with_pseudo_labels.pth', 'rb') as f:\n",
    "    inst = torch.load(f)\n",
    "inst = {ins['image_id']:ins['instances'] for ins in inst}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'image_id': '0016/4006608424_000013',\n",
       "  'bbox': [42, 102, 537, 420],\n",
       "  'category_id': 0,\n",
       "  'score': 0.7028999999999997,\n",
       "  'tid': 2},\n",
       " {'image_id': '0016/4006608424_000013',\n",
       "  'bbox': [178, 90, 426, 317],\n",
       "  'category_id': 0,\n",
       "  'score': 0.46229999999999916,\n",
       "  'tid': 9}]"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_frame_annots['0016/4006608424_000013']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22976\n",
      "28560\n",
      "nb of testing frames without detection results: 168\n",
      "['0016/4006608424_000012', '0017/2810112808_000018', '0017/2810112808_000019', '0017/2810112808_000020', '0017/2810112808_000021', '0017/2810112808_000024', '0017/2810112808_000025', '0017/2810112808_000027', '0017/2810112808_000028', '0017/2810112808_000029', '0017/2810112808_000030', '0017/2810112808_000031', '0017/2810112808_000033', '0017/2810112808_000034', '0017/2810112808_000068', '0017/2810112808_000069', '0017/2810112808_000070', '0017/2810112808_000071', '0017/2810112808_000072', '0017/2810112808_000073', '0017/2810112808_000074', '0017/2810112808_000076', '0017/2810112808_000077', '0017/2810112808_000078', '0019/4759861822_000012', '0028/5840177726_000011', '0050/11587211476_000041', '0050/11587211476_000042', '0050/11587211476_000046', '0050/11587211476_000047', '0074/7453733046_000024', '0074/7453733046_000032', '0074/7453733046_000033', '0080/5522880149_000024', '0080/5522880149_000025', '0080/5522880149_000026', '0080/5522880149_000027', '0081/6139126979_000005', '0082/11503803033_000016', '0082/11503803033_000018', '1000/2716277960_000006', '1000/2716277960_000012', '1000/4925211209_000006', '1000/4925211209_000007', '1000/4925211209_000008', '1000/4925211209_000009', '1000/4925211209_000010', '1000/4925211209_000011', '1000/4925211209_000012', '1000/4925211209_000013', '1000/4925211209_000014', '1000/4925211209_000015', '1000/4925211209_000016', '1000/4925211209_000017', '1000/4925211209_000018', '1000/4925211209_000019', '1000/6253433310_000019', '1000/6253433310_000020', '1000/6253433310_000021', '1000/6253433310_000022', '1000/6253433310_000023', '1001/4889681401_000042', '1001/4889681401_000043', '1001/4889681401_000044', '1001/4889681401_000045', '1001/4889681401_000046', '1001/4889681401_000047', '1001/4889681401_000048', '1001/6713120511_000002', '1001/6713120511_000003', '1001/6713120511_000004', '1001/6713120511_000005', '1002/2932897373_000007', '1002/4103088549_000006', '1002/4103088549_000007', '1002/4103088549_000008', '1002/4103088549_000009', '1005/5991060898_000021', '1005/5991060898_000022', '1005/5991060898_000023', '1008/5783819683_000002', '1008/5783819683_000003', '1008/8797589693_000007', '1008/8797589693_000031', '1008/8797589693_000032', '1008/8797589693_000033', '1008/8797589693_000051', '1008/8797589693_000080', '1011/8627074061_000002', '1011/8627074061_000003', '1011/8627074061_000004', '1011/8627074061_000060', '1011/8627074061_000061', '1015/5919180502_000038', '1015/5919180502_000039', '1015/5919180502_000040', '1015/5919180502_000041', '1015/5919180502_000042', '1015/5919180502_000043', '1021/5352022985_000003', '1021/5352022985_000004', '1021/5352022985_000007', '1021/5352022985_000008', '1021/5352022985_000009', '1021/5352022985_000022', '1021/5352022985_000024', '1021/5352022985_000025', '1021/5352022985_000026', '1025/8787109801_000024', '1027/3113970118_000003', '1051/5979720550_000005', '1051/5979720550_000006', '1051/5979720550_000007', '1051/5979720550_000008', '1103/2510696559_000041', '1103/3441428429_000021', '1103/3441428429_000022', '1103/3441428429_000023', '1103/3441428429_000024', '1103/3441428429_000025', '1103/3441428429_000026', '1103/3441428429_000027', '1103/3441428429_000031', '1103/3441428429_000032', '1103/3441428429_000033', '1103/3441428429_000036', '1103/3441428429_000045', '1103/3441428429_000048', '1103/3441428429_000049', '1103/3441428429_000050', '1103/3441428429_000051', '1103/3441428429_000052', '1103/3441428429_000053', '1103/3441428429_000072', '1103/3441428429_000073', '1103/5521781780_000036', '1104/2821968703_000029', '1110/3455926688_000002', '1110/3455926688_000003', '1110/3455926688_000004', '1110/3455926688_000010', '1110/3455926688_000011', '1110/3455926688_000012', '1110/3455926688_000013', '1110/3455926688_000014', '1110/3455926688_000016', '1110/3455926688_000017', '1110/3455926688_000020', '1110/3455926688_000021', '1110/3455926688_000022', '1110/3455926688_000023', '1110/3455926688_000024', '1124/13569831214_000002', '1124/13569831214_000003', '1124/13569831214_000004', '1124/13569831214_000009', '1124/13569831214_000033', '1150/6104044648_000002', '1150/6104044648_000004', '1150/6104044648_000010', '1201/5658916668_000016', '1201/5658916668_000017', '1201/5658916668_000018', '1201/5658916668_000019', '1201/5658916668_000022', '1203/3345608051_000054', '1203/3345608051_000055', '1203/3345608051_000056']\n"
     ]
    }
   ],
   "source": [
    "print(len(inst.keys()))\n",
    "print(len(val_frame_annots.keys()))\n",
    "subtract = set(inst.keys()) - set(val_frame_annots.keys())\n",
    "print('nb of testing frames without detection results:', len(subtract))\n",
    "print(sorted(subtract))\n",
    "# print(set(val_frame_annots.keys()) - set(inst.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "28560\n",
      "28560\n"
     ]
    }
   ],
   "source": [
    "print(len(val_frame_annots.keys()))\n",
    "print(len(set(val_frame_annots.keys())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'image_id': '0016/4006608424_000013',\n",
       "  'bbox': [42, 102, 537, 420],\n",
       "  'category_id': 0,\n",
       "  'score': 0.7028999999999997,\n",
       "  'tid': 2},\n",
       " {'image_id': '0016/4006608424_000013',\n",
       "  'bbox': [178, 90, 426, 317],\n",
       "  'category_id': 0,\n",
       "  'score': 0.46229999999999916,\n",
       "  'tid': 9}]"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_frame_annots['0016/4006608424_000013']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bboxs': [194, 21, 318, 160],\n",
       " 'class': 17,\n",
       " 'score': 0.8494999999999995,\n",
       " 'tid': 4}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# len(det_val_frame_annots) # 156725\n",
    "det_val_frame_annots[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# generate frame lists\n",
    "data = []\n",
    "i = 0\n",
    "annot_dir = 'annotation/training'\n",
    "for folder in tqdm(os.listdir(annot_dir)):\n",
    "    for video_json in os.listdir(os.path.join(annot_dir, folder)):\n",
    "        video_id = folder + '/' + video_json.split('.')[0]\n",
    "        if video_id not in train_video_dict:\n",
    "            continue\n",
    "        with open(os.path.join(annot_dir, folder, video_json), 'r') as f:\n",
    "            annot = json.load(f)\n",
    "        for j in range(annot['frame_count']):\n",
    "            data.append(\n",
    "                (folder + '/' + annot['video_id'], i, j, os.path.join(folder, annot['video_id'], annot['video_id'] + '_' + str(f'{j + 1:06d}') + '.jpg'), '')\n",
    "            )\n",
    "        i += 1\n",
    "    \n",
    "train_frame_list = pd.DataFrame(data, columns = ['original_video_id' , 'video_id', 'frame_id' , 'path', 'labels'])\n",
    "train_frame_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "data = []\n",
    "i = 0\n",
    "annot_dir = 'annotation/validation'\n",
    "for folder in tqdm(os.listdir(annot_dir)):\n",
    "    for video_json in os.listdir(os.path.join(annot_dir, folder)):\n",
    "        video_id = folder + '/' + video_json.split('.')[0]\n",
    "        if video_id not in val_video_dict:\n",
    "            continue\n",
    "        with open(os.path.join(annot_dir, folder, video_json), 'r') as f:\n",
    "            annot = json.load(f)\n",
    "        for j in range(annot['frame_count']):\n",
    "            data.append(\n",
    "                (folder + '/' + annot['video_id'], i, j, os.path.join(folder, annot['video_id'], annot['video_id'] + '_' + str(f'{j + 1:06d}') + '.jpg'), '')\n",
    "            )\n",
    "        i += 1\n",
    "    \n",
    "val_frame_list = pd.DataFrame(data, columns = ['original_video_id' , 'video_id', 'frame_id' , 'path', 'labels'])\n",
    "val_frame_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_frame_list.to_csv('frame_lists/train.csv', sep=' ', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_frame_list.to_csv('frame_lists/val.csv', sep=' ', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "slowfast",
   "language": "python",
   "name": "slowfast"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
